import argparse
import os

from combination_train.resnet_combination_train import resnet_combination_train
from combination_train.bnnt_train import bnnt_train_cifar10, bnnt_train_mnist
from util.wordNetCifar10 import *
from util.utils import mkdirs

train_model_dict = {
    0: resnet_combination_train,
    1: bnnt_train_cifar10,
    2: bnnt_train_mnist,
}

def show_dict(d):
    [print(e) for e in [f"{key}: {d[key].__name__}" for key in d]]

arg_cli = argparse.ArgumentParser()
arg_cli.add_argument("--train_number", help=f"训练模型编号", type=int)
arg_cli.add_argument("--wordnet", help="显示wordnet-cifar10上位图", action="store_true")
# arg_cli.add_argument("--bnn_tree", help="显示bnn-tree", action="store_true")
args = arg_cli.parse_args()

if __name__ == "__main__":
    # resnet_combination_train()
    # bnnt_train_cifar10()
    # bnnt_train_mnist()
    if args.train_number in train_model_dict:
        train_model_dict[args.train_number]
    else:
        show_dict(train_model_dict)

    if args.wordnet:
        generate_tree()
        dir_ = os.path.join(*["output", "wordnet-cifar10"])
        mkdirs(os.path.join(dir_))
        digraph.view(directory=dir_)

    # if args.bnn_tree:
    #     from util.bnn_tree_visualize import *
    #     from NN.bnn_tree import *
    #     from util.bnnt_dataset import BNNTDATA
    #     bnntree = BNNTree(10, 1, None, None, 0)
    #     draw_tree(bnntree.tree)

    